from setting import data
import polars as pl
import pickle
def test_data():
    # 获取回测数据
    base_data, trade_data = data.get_data(
        ["SHFE_ag00_60"],                          # 基准合约列表
        ["CFFEX_IF00_60", "CFFEX_IC00_60"],         # 交易合约池列表  
        "20250101",                                 # 开始日期
        "20251001"                                  # 结束日期
    )
    print(base_data)
    base_data = base_data.with_columns([
        # 判断是否是周一早上10点
        ((pl.col("datetime").dt.weekday() == 1) &  # 周一 = 1
         (pl.col("datetime").dt.hour() == 10) &    # 10点
         (pl.col("datetime").dt.minute() == 0))    # 0分
        .cast(pl.Int8).alias("buy"),
        
        # 判断是否是周五下午2点
        ((pl.col("datetime").dt.weekday() == 5) &  # 周五 = 5
         (pl.col("datetime").dt.hour() == 14) &    # 14点
         (pl.col("datetime").dt.minute() == 0))    # 0分
        .cast(pl.Int8).alias("sell")
    ])
    for row in base_data.iter_rows():
        print(row) 

def test_parms():
    with open("setting/volume_multiple.pkl", "rb") as f:
        loaded_volume = pickle.load(f)
    with open("setting/margin_ratio.pkl", "rb") as f:
        loaded_margin = pickle.load(f)
    deposit_rates = loaded_margin
    multipliers = loaded_volume
    print(deposit_rates)
    print(multipliers)

def test_load():
    trade_data_symbol = [
        "CZCE_AP00_60", "CZCE_CF00_60", "CZCE_CJ00_60", "CZCE_CY00_60", "CZCE_FG00_60", 
        "CZCE_LR00_60", "CZCE_MA00_60", "CZCE_OI00_60", "CZCE_PF00_60", "CZCE_PK00_60", 
        "CZCE_PL00_60", "CZCE_PM00_60", "CZCE_PR00_60", "CZCE_RI00_60", "CZCE_RM00_60", 
        "CZCE_RS00_60", "CZCE_SA00_60", "CZCE_SH00_60", "CZCE_SF00_60", "CZCE_SM00_60", 
        "CZCE_SR00_60", "CZCE_TA00_60", "CZCE_UR00_60", "CZCE_WH00_60", "CZCE_ZC00_60", 
        "DCE_a00_60", "DCE_b00_60", "DCE_bb00_60", "DCE_bz00_60", "DCE_c00_60", 
        "DCE_cs00_60", "DCE_eb00_60", "DCE_eg00_60", "DCE_fb00_60", "DCE_j00_60", 
        "DCE_jd00_60", "DCE_jm00_60", "DCE_l00_60", "DCE_lg00_60", "DCE_lh00_60", 
        "DCE_m00_60", "DCE_p00_60", "DCE_pg00_60", "DCE_pp00_60", "DCE_rr00_60", 
        "DCE_v00_60", "DCE_y00_60", "GFEX_lc00_60", "GFEX_ps00_60", "GFEX_si00_60", 
        "INE_bc00_60", "INE_ec00_60", "INE_lu00_60", "INE_nr00_60", "INE_sc00_60", 
        "SHFE_ad00_60", "SHFE_ag00_60", "SHFE_al00_60", "SHFE_ao00_60", "SHFE_br00_60", 
        "SHFE_bu00_60", "SHFE_cu00_60", "SHFE_fu00_60", "SHFE_hc00_60", "SHFE_ni00_60", 
        "SHFE_op00_60", "SHFE_pb00_60", "SHFE_rb00_60", "SHFE_ru00_60", "SHFE_sn00_60", 
        "SHFE_sp00_60", "SHFE_ss00_60", "SHFE_wr00_60", "SHFE_zn00_60"
    ]
    base_data, all_trade_data = data.get_data(["SHFE_ag00_60"], trade_data_symbol, "20250101", "20251001")
    print(base_data)
if __name__ == "__main__":
    # test_parms()
    # test_data()
    test_load()